# -*- coding:utf-8 -*-
from fastapi import FastAPI, UploadFile, File
from fastapi.middleware.cors import CORSMiddleware
from starlette.responses import FileResponse
import psutil
import speedtest
import os
import shutil
import uvicorn
import zipfile
import time
app = FastAPI()

origins = [
    "http://localhost",
    "http://localhost:8000",
    "http://localhost:8100",
    "http://localhost:8181",

    "http://127.0.0.1",
    "http://127.0.0.1:8000",
    "http://127.0.0.1:8100",
    "http://127.0.0.1:8181",

    "http://120.53.88.254",
    "http://120.53.88.254:8000",
    "http://120.53.88.254:8100",
    "http://120.53.88.254:8181",

    "http://82.157.104.42",
    "http://82.157.104.42:8000",
    "http://82.157.104.42:8100",
    "http://82.157.104.42:8181",
]

app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

'''
启动
pip install fastapi psutil speedtest-cli "uvicorn[standard]" python-multipart
uvicorn main:app --reload  --host 0.0.0.0 --port 7079
'''


'''
虚拟机网速测试
'''
@app.get("/test-internet-speed")
async def test_internet_speed():
    try:
        test = speedtest.Speedtest()
        test.get_servers()
        test.get_best_server()
        download_speed = float(test.download() / 1024)
        upload_speed = float(test.upload() / 1024)
    except Exception as e:
        #print("error s" % e)
        return {
            "code": 201,
            "data": {}
        }
    else:
        return {
            "code": 200,
            "data": {
                'download_speed': round(download_speed, 2),
                'upload_speed': round(upload_speed, 2),
            }
        }

'''
虚拟机磁盘空间详情
'''
@app.get("/df-h")
async def dh_f():
    try:
        partitions = psutil.disk_partitions()
        df = []
        for partition in partitions:
            if not os.path.exists(partition.mountpoint):
                continue

            usage = psutil.disk_usage(partition.mountpoint)
            df.append({
                'device': partition.device,
                'total': round(usage.total / (1024 ** 3), 2),
                'used': round(usage.used / (1024 ** 3), 2),
                'free': round(usage.free / (1024 ** 3), 2)
            })
    except Exception as e:
        #print("error s" % e)
        return {
            "code": 201,
            "data": {},
        }
    else:
        return {
            "code": 200,
            "data": {
                'df_h': df,
                'total': {
                    'device': "&".join(list(map(lambda x: x['device'], df))),
                    'total': round(sum(list(map(lambda x: x['total'], df))), 2),
                    'used': round(sum(list(map(lambda x: x['used'], df))), 2),
                    'free': round(sum(list(map(lambda x: x['free'], df))), 2)
                }
            }
        }

'''
虚拟机桌面文件中转站中的文件列表
'''
@app.get("/filetransfer-station-files")
async def get_file_transfer_station_files():
    _, file_transfer_station_files = _get_desktop_path()
    return {
        "code": 200,
        "data": {
            'file_transfer_station_files': file_transfer_station_files
        }
    }


'''
从虚拟机桌面文件中转站中下载文件
'''
@app.get("/file-download")
async def file_download(fileName: str):
    files = fileName.split(",")
    filePath, zip_name = _zip_files(files)
    return FileResponse(
        filePath,
        **{"filename": f"{zip_name}"}
    )

'''
上传文件到虚拟机桌面文件中转站
'''
@app.post("/file-upload")
async def file_upload(file: UploadFile =  File(...)):
    file_transfer_station_path, _ = _get_desktop_path()
    file_path = os.path.join(file_transfer_station_path, file.filename)
    with open(file_path, "wb") as f:
        shutil.copyfileobj(file.file, f)
    return {
        "code": 200,
        "data": {
            "fileName": file.filename,
            "filePath": file_path,
            "contentType": file.content_type
        }
    }

def _get_desktop_path(stationName='文件中转站', file_formats=['.txt', '.doc', '.docx', '.xls', '.pdf']):   # in ubuntu is FileTransferStation
    desktop_path = os.path.join(os.path.expanduser("~"), 'Desktop')
    file_transfer_station_path = os.path.join(desktop_path, stationName)
    file_transfer_station_files = [
        file for file in os.listdir(file_transfer_station_path)
        if os.path.isfile(os.path.join(file_transfer_station_path, file))
        and any(file.endswith(file_format) for file_format in file_formats)
    ]
    return file_transfer_station_path, file_transfer_station_files

# def _get_desktop_path(stationName = '文件中转站'):             # in ubuntu is FileTransferStation
#     desktop_path = os.path.join(os.path.expanduser("~"), 'Desktop')
#     file_transfer_station_path = os.path.join(desktop_path, stationName)
#     file_transfer_station_files = [file for file in os.listdir(file_transfer_station_path) if os.path.isfile(os.path.join(file_transfer_station_path, file))]
#     return file_transfer_station_path, file_transfer_station_files

# def _zip_files(files):
#     file_transfer_station_path, _ = _get_desktop_path()
#     total_original_size = 0  # 初始化原文件大小总和为0
#     if len(files) == 1:
#         file_path = os.path.join(file_transfer_station_path, files[0])
#         return file_path, files[0], os.path.getsize(file_path) / (1024 * 1024)  # 返回文件路径、文件名和大小（MB）
#     files_to_compress = [os.path.join(file_transfer_station_path, file) for file in files]
#     zip_name = f"{time.time()}.zip"
#     zip_path = os.path.join(file_transfer_station_path, zip_name)
#     with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zip_ref:
#         for file_name in files_to_compress:
#             original_file_size = os.path.getsize(file_name)
#             total_original_size += original_file_size  # 累加原文件大小
#             zip_ref.write(file_name, arcname=file_name.replace(file_transfer_station_path, "files"))
#     total_original_size_mb = total_original_size / (1024 * 1024)  # 将原文件大小之和转换为MB
#     return zip_path, zip_name, total_original_size_mb

def _zip_files(files):
    file_transfer_station_path, _ = _get_desktop_path()
    if len(files) == 1:
        return os.path.join(file_transfer_station_path, files[0]), files[0]
    files_to_compress = [os.path.join(file_transfer_station_path, file) for file in files]
    zip_name = f"{time.time()}.zip"
    zip_path = os.path.join(file_transfer_station_path, zip_name)
    with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zip_ref:
        for file_name in files_to_compress:
            zip_ref.write(file_name, **{"arcname": file_name.replace(file_transfer_station_path, "files" )})
    return zip_path, zip_name


# in ubuntu use, in windows use start.bat
if __name__ == '__main__':
    uvicorn.run('main:app', host='0.0.0.0', port=7079, log_level="info")


